import sympy as sp
from sympy.physics.hydrogen import *
from manimlib import *
from scipy import optimize


def Wf_nlm(n, l, m):
    x = sp.Symbol("x", real=True, positive=True)
    y = sp.Symbol("y", real=True, positive=True)
    z = sp.Symbol("z", real=True, positive=True)
    r = sp.sqrt(x**2+y**2+z**2)
    phi = sp.atan2(y, x)
    theta = sp.acos(z/r)

    Wave = Psi_nlm(n, l, m, r, phi, theta, 1)
    Prob = Wave * sp.conjugate(Wave)
    Prob = sp.simplify(Prob)

    R = sp.Symbol("R", real=True, positive=True)
    Rad_Wave = sp.simplify(R_nl(n, l, R, 1)**2)
    Theta = sp.Symbol("Theta")
    Tangent_Wave = sp.simplify(sp.Ynm(l, m, Theta, 0).expand(func=True)**2)

    f_Prob = sp.lambdify([x, y, z], Prob, 'numpy')
    f_Rad = sp.lambdify(R, -Rad_Wave, 'numpy')
    f_Tan = sp.lambdify(Theta, -Tangent_Wave, 'numpy')

    maxi1 = optimize.fminbound(f_Rad, 0, 10)
    maxi2 = optimize.fminbound(f_Tan, 0, PI)
    return (f_Prob, f_Rad(maxi1)*f_Tan(maxi2))


def random_posi(r, aprox_size=10000):
    p = np.random.uniform(-r, r, (int(aprox_size*6//PI), 3))
    p = np.delete(p, np.where(np.linalg.norm(p, axis=1) > r), axis=0)
    return p


def alternating(n):
    inital = np.arange(0,2*n+1)
    return (inital+1) // 2 * (-1)**(inital+1)
print(alternating(4))



class HydrogenCloud(DotCloud):
    def __init__(self, n, l, m, sample=500, r=2, **kwargs):
        print(f"Drawing Orbit {n}{l}{m}...", end="")
        self.Wf, self.stre = Wf_nlm(n, l, m)
        self.r = r*n**2
        self.n = n
        self.l = l
        self.m = m
        coords = self.get_randpoints(sample*10)
        while len(coords) < sample:
            coords = np.append(coords, self.get_randpoints(sample*10), axis=0)
        coords = coords[0:sample]
        super().__init__(coords, **kwargs)
        self.make_3d()
        self.set_color_by_rgba_func(self.color_func)
        self.scale(4/self.r, False, about_point=ORIGIN)
        print("Finished.")

    def color_func(self, p):
        prob = self.Wf(p[0], p[1], p[2])/self.stre
        func = get_rgb_gradient_function(-3, 0, '3b1b_colormap')
        return np.array([*func(np.log10(prob)), self.opacity])

    def Siever(self, points):
        x = points[:, 0]
        y = points[:, 1]
        z = points[:, 2]
        data = self.Wf(x, y, z)
        rand_data = np.random.uniform(0, self.stre, x.shape)
        return np.delete(points, np.where(rand_data > data), axis=0)

    def get_randpoints(self, sample):
        init_points = random_posi(self.r, sample*10)
        return self.Siever(init_points)

    def get_tex(self):
        formula = Psi_nlm(self.n, self.l, self.m, sp.Symbol(
            "r"), sp.Symbol("varphi"), sp.Symbol("theta"))
        formula = sp.simplify(formula)
        tex = Tex(
            "\Psi_{"+str(self.n)+","+str(self.l)+"," +
            str(self.m)+r"}(r,\varphi,\theta)=",
            sp.printing.latex(formula),
        )
        return tex

    def get_orbname(self):
        l_map = ["s", "p", "d", "f", "g", "h", "i"]
        name = f"{self.n}\mathrm {l_map[self.l]}_"+"{"+str(self.m)+"}"
        return Tex(name)


class TestRand(Scene):
    def construct(self):
        axis = ThreeDAxes([-3, 3], [-3, 3], [-3, 3],
                          axis_config={"including_tips": True}, opacity=0.2).apply_depth_test()
        self.add(axis)
        frame = self.camera.frame
        frame.set_euler_angles(30*DEGREES, 70*DEGREES)
        frame.add_updater(lambda m: m.set_euler_angles(
            self.time/10, PI/4+PI/6*np.sin(self.time/4)))
        self.add(frame)
        H = HydrogenCloud(1, 0, 0, r=4, radius=0.01, opacity=0.8, sample=50000)
        self.add(H)

        l = Line(
            DOWN*3, UP*3).insert_n_curves(50).set_color(COLORMAP_3B1B).set_stroke(width=10)
        prob_hi = TexText("高概率").scale(0.7).next_to(
            l, UP).set_backstroke(width=5)
        prob_lo = TexText("低概率").scale(0.7).next_to(
            l, DOWN).set_backstroke(width=5)
        tex = H.get_tex().fix_in_frame().scale(0.7).to_edge(DOWN)
        tex.set_backstroke(width=5)
        vg = VGroup(l, prob_hi, prob_lo).fix_in_frame().to_edge(RIGHT, buff=1)
        self.add(vg, tex)
        name = H.get_orbname()
        name.set_backstroke(width=5)
        name_vg = VGroup(name, TexText("轨道")).arrange(
            DOWN).scale(2).to_edge(LEFT, buff=1).fix_in_frame()
        self.add(name_vg)
        for n in range(2, 8):
            for l in range(0, n):
                for m in alternating(l):
                    H2 = HydrogenCloud(
                        n, l, m, r=4, radius=0.01, opacity=0.8, sample=50000)
                    tex2 = H2.get_tex().fix_in_frame().scale(0.7).to_edge(DOWN)
                    name2 = H2.get_orbname().fix_in_frame().scale(2).move_to(name)
                    tex2.set_backstroke(width=5)
                    name2.set_backstroke(width=5)
                    self.wait(5)
                    self.play(
                        Transform(H, H2),
                        FadeTransform(
                            tex, tex2, fade_transform_mismatches=True),
                        FadeTransform(name, name2)
                    )
                    self.remove(tex2, name2)
                    tex.become(tex2)
                    name.become(name2)
                    self.add(tex, name)
        self.wait(5)

class HydrogenScene(Scene):
    CONFIG = {
        "n": 7,
        "l": 2,
        "m": 2,
    }
    def construct(self):
        axis = ThreeDAxes([-3, 3], [-3, 3], [-3, 3],
                          axis_config={"including_tips": True}, opacity=0.2).apply_depth_test()
        self.add(axis)
        frame = self.camera.frame
        frame.set_euler_angles(30*DEGREES, 70*DEGREES)
        # frame.add_updater(lambda m: m.set_euler_angles(
        #     self.time/10, PI/4+PI/4*np.sin(self.time/4)))
        self.add(frame)
        H = HydrogenCloud(self.n, self.l, self.m, r=4, radius=0.01, opacity=0.8, sample=50000)
        self.add(H)

        l = Line(
            DOWN*3, UP*3).insert_n_curves(50).set_color(COLORMAP_3B1B).set_stroke(width=10)
        prob_hi = TexText("高概率").scale(0.7).next_to(
            l, UP).set_backstroke(width=5)
        prob_lo = TexText("低概率").scale(0.7).next_to(
            l, DOWN).set_backstroke(width=5)
        tex = H.get_tex().fix_in_frame().scale(0.7).to_edge(DOWN)
        tex.set_max_width(10)
        tex.set_backstroke(width=5)
        vg = VGroup(l, prob_hi, prob_lo).fix_in_frame().to_edge(RIGHT, buff=1)
        self.add(vg, tex)
        name = H.get_orbname()
        name.set_backstroke(width=5)
        name_vg = VGroup(name, TexText("轨道")).arrange(
            DOWN).scale(2).to_edge(LEFT, buff=1).fix_in_frame()
        self.add(name_vg)


class Thumbnail(Scene):
    def construct(self):
        axis = ThreeDAxes([-3, 3], [-3, 3], [-3, 3],
                          axis_config={"including_tips": True}, opacity=0.2).apply_depth_test()
        self.add(axis)
        frame = self.camera.frame
        frame.set_euler_angles(30*DEGREES, 20*DEGREES)
        # # frame.add_updater(lambda m: m.set_euler_angles(
        # #     self.time/10, PI/4+PI/4*np.sin(self.time/4)))
        # self.add(frame)
        H = HydrogenCloud(7, 2, 2, r=4, radius=0.01, opacity=.95, sample=50000)
        self.add(H)    
        # cir = Circle(radius=2.5)
        # self.add(cir)

        text = TexText("氢原子波函数","可视化").set_bstroke(width=10).fix_in_frame().arrange(RIGHT,buff=2)
        text.rotate(-PI/2)
        text[1].rotate(PI)
        text.set_height(2*PI-1.5)
        text[0].move_to(PI/2*UP)
        text[1].move_to(-PI/2*UP)
        text.shift(RIGHT*np.log(3))


        text.apply_complex_function(lambda z: np.exp(z))

        self.add(text)

        self.add(Tex("7\mathrm p_2").scale(1.5).set_bstroke(width=10).fix_in_frame())